#!/usr/bin/env python3
"""
Database initialization script for EmailQA system
This script reads SQL files from the sql/ directory and executes them in order
"""
import os
import sys
import sqlite3
import pymysql
from pathlib import Path
from dotenv import load_dotenv

# Add src to path
sys.path.append('src')

load_dotenv()

class DatabaseInitializer:
    """Database initialization class"""
    
    def __init__(self, host=None, port=None, user=None, password=None, database=None):
        self.host = host or os.getenv('MYSQL_HOST', 'localhost')
        self.port = port or int(os.getenv('MYSQL_PORT', 10000))
        self.user = user or os.getenv('MYSQL_USER', 'root')
        self.password = password or os.getenv('MYSQL_PASSWORD', '123456')
        self.database = database or os.getenv('MYSQL_DATABASE', 'emailqa')
        self.sql_dir = Path(__file__).parent / 'sql'
        self.mysql_dir = self.sql_dir / 'mysql'
        self.sqlite_dir = self.sql_dir / 'sqlite'
        
    def get_connection(self):
        """Get a MySQL database connection"""
        try:
            connection = pymysql.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database,
                charset='utf8mb4',
                cursorclass=pymysql.cursors.DictCursor
            )
            return connection
        except pymysql.Error as e:
            raise Exception(f"Failed to connect to MySQL: {e}")
    
    def test_connection(self) -> bool:
        """Test database connection"""
        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute("SELECT 1")
                    return True
        except Exception:
            return False
    
    def get_mysql_sql_files(self):
        """Get sorted list of MySQL SQL files"""
        if not self.mysql_dir.exists():
            print(f"MySQL directory not found: {self.mysql_dir}")
            return []
        
        # Get all .sql files and sort by filename
        sql_files = [f for f in self.mysql_dir.glob('*.sql')]
        sql_files.sort(key=lambda x: x.name)
        
        return sql_files
    
    def get_sqlite_sql_files(self):
        """Get sorted list of SQLite SQL files"""
        if not self.sqlite_dir.exists():
            print(f"SQLite directory not found: {self.sqlite_dir}")
            return []
        
        # Get all .sql files and sort by filename
        sql_files = [f for f in self.sqlite_dir.glob('*.sql')]
        sql_files.sort(key=lambda x: x.name)
        
        return sql_files
    
    def execute_sql_file(self, file_path):
        """Execute a single SQL file"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                sql_content = f.read()
            
            # Remove comments and empty lines
            lines = [line.strip() for line in sql_content.split('\n') 
                     if line.strip() and not line.strip().startswith('--')]
            sql_content = '\n'.join(lines)
            
            # Split into individual statements
            statements = [stmt.strip() for stmt in sql_content.split(';') if stmt.strip()]
            
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    for statement in statements:
                        if statement:
                            cursor.execute(statement)
                    conn.commit()
                    
            print(f"[OK] Executed: {file_path.name}")
            return True
            
        except Exception as e:
            print(f"[FAIL] Error executing {file_path.name}: {e}")
            return False
    
    def initialize_mysql_database(self):
        """Initialize MySQL database with MySQL SQL files"""
        print("Initializing MySQL Database...")
        print("=" * 40)
        
        # Test connection first
        if not self.test_connection():
            print("[FAIL] Cannot connect to MySQL database")
            return False
        
        print(f"[OK] Connected to MySQL at {self.host}:{self.port}/{self.database}")
        
        # Get MySQL SQL files
        sql_files = self.get_mysql_sql_files()
        if not sql_files:
            print("[WARNING] No MySQL SQL files found")
            return False
        
        print(f"[INFO] Found {len(sql_files)} MySQL SQL files")
        
        # Execute each SQL file
        success_count = 0
        for sql_file in sql_files:
            if self.execute_sql_file(sql_file):
                success_count += 1
        
        print("=" * 40)
        if success_count == len(sql_files):
            print(f"[SUCCESS] MySQL database initialized successfully ({success_count}/{len(sql_files)} files)")
            return True
        else:
            print(f"[PARTIAL] MySQL database partially initialized ({success_count}/{len(sql_files)} files)")
            return False
    
    def initialize_sqlite_database(self, db_path='emailqa.db'):
        """Initialize SQLite database with SQLite SQL files"""
        print("Initializing SQLite Database...")
        print("=" * 40)
        
        # Get SQLite SQL files
        sql_files = self.get_sqlite_sql_files()
        if not sql_files:
            print("[WARNING] No SQLite SQL files found")
            return False
        
        print(f"[INFO] Found {len(sql_files)} SQLite SQL files")
        print(f"[INFO] SQLite database path: {db_path}")
        
        # Execute each SQL file
        success_count = 0
        for sql_file in sql_files:
            if self.execute_sqlite_sql_file(sql_file, db_path):
                success_count += 1
        
        print("=" * 40)
        if success_count == len(sql_files):
            print(f"[SUCCESS] SQLite database initialized successfully ({success_count}/{len(sql_files)} files)")
            return True
        else:
            print(f"[PARTIAL] SQLite database partially initialized ({success_count}/{len(sql_files)} files)")
            return False
    
    def execute_sqlite_sql_file(self, file_path, db_path):
        """Execute a single SQLite SQL file"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                sql_content = f.read()
            
            # Remove comments and empty lines
            lines = [line.strip() for line in sql_content.split('\n') 
                     if line.strip() and not line.strip().startswith('--')]
            sql_content = '\n'.join(lines)
            
            # Split into individual statements
            statements = [stmt.strip() for stmt in sql_content.split(';') if stmt.strip()]
            
            with sqlite3.connect(db_path) as conn:
                cursor = conn.cursor()
                for statement in statements:
                    if statement:
                        cursor.execute(statement)
                conn.commit()
                    
            print(f"[OK] Executed SQLite: {file_path.name}")
            return True
            
        except Exception as e:
            print(f"[FAIL] Error executing SQLite {file_path.name}: {e}")
            return False
    
    def initialize_database(self):
        """Initialize both MySQL and SQLite databases"""
        print("Initializing EmailQA Databases...")
        print("=" * 50)
        
        mysql_success = self.initialize_mysql_database()
        sqlite_success = self.initialize_sqlite_database()
        
        print("=" * 50)
        if mysql_success and sqlite_success:
            print("[SUCCESS] Both MySQL and SQLite databases initialized successfully!")
            return True
        elif mysql_success or sqlite_success:
            print("[PARTIAL] Database initialization partially successful")
            return False
        else:
            print("[FAILED] Both MySQL and SQLite database initialization failed!")
            return False
    
    def show_mysql_database_info(self):
        """Show MySQL database information"""
        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    # Get table list
                    cursor.execute("""
                        SELECT table_name, table_comment 
                        FROM information_schema.tables 
                        WHERE table_schema = %s
                        ORDER BY table_name
                    """, (self.database,))
                    tables = cursor.fetchall()
                    
                    print(f"\nMySQL Database: {self.database}")
                    print("Tables:")
                    for table in tables:
                        print(f"  - {table['table_name']}: {table['table_comment'] or 'No comment'}")
                        
        except Exception as e:
            print(f"Error getting MySQL database info: {e}")
    
    def show_sqlite_database_info(self, db_path='emailqa.db'):
        """Show SQLite database information"""
        try:
            if not os.path.exists(db_path):
                print(f"\nSQLite database file not found: {db_path}")
                return
            
            with sqlite3.connect(db_path) as conn:
                cursor = conn.cursor()
                
                # Get table list
                cursor.execute("""
                    SELECT name, sql 
                    FROM sqlite_master 
                    WHERE type='table' 
                    ORDER BY name
                """)
                tables = cursor.fetchall()
                
                print(f"\nSQLite Database: {db_path}")
                print("Tables:")
                for table in tables:
                    print(f"  - {table[0]}")
                        
        except Exception as e:
            print(f"Error getting SQLite database info: {e}")
    
    def show_database_info(self):
        """Show both MySQL and SQLite database information"""
        self.show_mysql_database_info()
        self.show_sqlite_database_info()

def main():
    """Main function"""
    print("EmailQA Database Initialization Script")
    print("=" * 50)
    
    try:
        initializer = DatabaseInitializer()
        
        # Show current database info
        initializer.show_database_info()
        
        # Initialize database
        if initializer.initialize_database():
            print("\n[SUCCESS] Database initialization completed!")
            
            # Show updated database info
            initializer.show_database_info()
            
        else:
            print("\n[FAILED] Database initialization failed!")
            sys.exit(1)
            
    except KeyboardInterrupt:
        print("\n[INFO] Operation cancelled by user")
        sys.exit(0)
    except Exception as e:
        print(f"\n[ERROR] {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()